import torch.nn.functional as F
import torch
import torch.nn as nn
from torch.nn.utils import weight_norm


class Chomp1d(nn.Module):
    def __init__(self, chomp_size):
        super(Chomp1d, self).__init__()
        self.chomp_size = chomp_size

    def forward(self, x):
        return x[:, :, :-self.chomp_size].contiguous()


class TemporalBlock(nn.Module):
    def __init__(self, n_inputs, n_outputs, kernel_size, stride, dilation, padding, dropout=0.2):
        super(TemporalBlock, self).__init__()
        self.conv1 = weight_norm(nn.Conv1d(n_inputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))
        self.chomp1 = Chomp1d(padding)
        self.relu1 = nn.ReLU()
        self.dropout1 = nn.Dropout(dropout)

        self.conv2 = weight_norm(nn.Conv1d(n_outputs, n_outputs, kernel_size,
                                           stride=stride, padding=padding, dilation=dilation))
        self.chomp2 = Chomp1d(padding)
        self.relu2 = nn.ReLU()
        self.dropout2 = nn.Dropout(dropout)

        self.net = nn.Sequential(self.conv1, self.chomp1, self.relu1, self.dropout1,
                                 self.conv2, self.chomp2, self.relu2, self.dropout2)
        self.downsample = nn.Conv1d(
            n_inputs, n_outputs, 1) if n_inputs != n_outputs else None
        self.relu = nn.ReLU()
        self.init_weights()

    def init_weights(self):
        self.conv1.weight.data.normal_(0, 0.01)
        self.conv2.weight.data.normal_(0, 0.01)
        if self.downsample is not None:
            self.downsample.weight.data.normal_(0, 0.01)

    def forward(self, x):
        out = self.net(x)
        res = x if self.downsample is None else self.downsample(x)
        return self.relu(out + res)


class TemporalConvNet(nn.Module):
    def __init__(self, num_inputs, num_channels, kernel_size=2, dropout=0.2):
        super(TemporalConvNet, self).__init__()
        layers = []
        num_levels = len(num_channels)
        for i in range(num_levels):
            dilation_size = 2 ** i
            in_channels = num_inputs if i == 0 else num_channels[i-1]
            out_channels = num_channels[i]
            layers += [TemporalBlock(in_channels, out_channels, kernel_size, stride=1, dilation=dilation_size,
                                     padding=(kernel_size-1) * dilation_size, dropout=dropout)]

        self.network = nn.Sequential(*layers)

    def forward(self, x):
        return self.network(x)


class AddEps(nn.Module):
    def __init__(self, channels):
        super(AddEps, self).__init__()

        self.channels = channels
        self.linear = nn.Linear(channels, channels)

    def forward(self, x):
        eps = torch.randn_like(x)
        eps = self.linear(eps)

        return eps + x


class Encoder(nn.Module):
    def __init__(self, hyper_params):
        super(Encoder, self).__init__()
        self.hyper_params = hyper_params

        self.tcn = TemporalConvNet(
            num_inputs=hyper_params['item_embed_size'],
            num_channels=[hyper_params['hidden_size'],
                          hyper_params['hidden_size']],
            kernel_size=2,
            dropout=0.2
        )

    def forward(self, x):
        x = self.tcn(x.transpose(1, 2))
        return x.transpose(1, 2)


class Decoder(nn.Module):
    def __init__(self, hyper_params):
        super(Decoder, self).__init__()
        self.linear1 = nn.Linear(
            hyper_params['latent_size'], hyper_params['item_embed_size'])
        self.linear2 = nn.Linear(
            hyper_params['item_embed_size'], hyper_params['total_items'] + 1)
        nn.init.xavier_normal_(self.linear1.weight)
        nn.init.xavier_normal_(self.linear2.weight)
        self.activation = nn.Tanh()

    def forward(self, x):
        x = self.linear1(x)
        out_embed = x

        x = self.activation(x)
        x = self.linear2(x)
        return x, out_embed


class Model(nn.Module):
    def __init__(self, hyper_params):
        super(Model, self).__init__()
        self.hyper_params = hyper_params

        self.encoder = Encoder(hyper_params)
        self.decoder = Decoder(hyper_params)

        self.item_embed = nn.Embedding(
            hyper_params['total_items'] + 1, hyper_params['item_embed_size'], padding_idx=0)

        self.linear_o = nn.Linear(
            hyper_params['hidden_size'], hyper_params['latent_size'])
        self.linear1 = nn.Linear(
            hyper_params['hidden_size'], 2 * hyper_params['latent_size'])
        nn.init.xavier_normal_(self.linear1.weight)

        self.tanh = nn.Tanh()
        self.embed_dropout = nn.Dropout(0.2)
        self.encoded_dropout = nn.Dropout(0.2)

        self.add_eps1 = AddEps(hyper_params['item_embed_size'])
        self.add_eps2 = AddEps(hyper_params['hidden_size'])

    def sample_latent(self, z_inferred):
        return torch.randn_like(z_inferred)

    def forward(self, x):
        x = self.item_embed(x)
        x_real = x
        # x = self.add_eps1(x)
        # z_inferred = self.encoder(self.add_eps1(x))
        z_inferred = self.encoder(x)
        # z_inferred = self.linear_o(self.add_eps2(z_inferred))
        z_inferred = self.linear_o(z_inferred)
        z_inferred = self.encoded_dropout(z_inferred)
        # [batch_size x seq_len x total_items]
        dec_out, out_embed = self.decoder(z_inferred)

        return dec_out, x_real, z_inferred, out_embed


class Embed(nn.Module):
    def __init__(self, hyper_params):
        super(Embed, self).__init__()

        self.item_embed = nn.Embedding(
            hyper_params['total_items'] + 1, hyper_params['item_embed_size'])

    def forward(self, x):
        return self.item_embed(x)


class Adversary(nn.Module):
    def __init__(self, hyper_params):
        super(Adversary, self).__init__()
        self.hyper_params = hyper_params
        self.linear_i = nn.Linear(
            hyper_params['item_embed_size'] + hyper_params['latent_size'], 128)

        self.dnet_list = []
        self.net_list = []
        for _ in range(2):
            self.dnet_list.append(nn.Linear(128, 128))
            self.net_list.append(nn.Linear(128, 128))

        self.dnet_list = nn.ModuleList(self.dnet_list)
        self.net_list = nn.ModuleList(self.net_list)

        self.linear_o = nn.Linear(128, hyper_params['latent_size'])
        self.dropout1 = nn.Dropout(0.2)
        self.dropout2 = nn.Dropout(0.2)

    def forward(self, x, z, padding):
        # batch_size x seq_len x dim
        net = torch.cat((x, z), 2)
        net = self.linear_i(net)
        net = self.dropout1(net)

        for i in range(2):
            dnet = self.dnet_list[i](net)
            net = net + self.net_list[i](dnet)
            net = F.elu(net)

        # seq_len
        net = self.linear_o(net)
        net = self.dropout2(net)
        net = net + 0.5 * torch.square(z)

        net = net * (1.0 - padding.float().unsqueeze(2))

        return net


class GRUAdversary(nn.Module):
    def __init__(self, hyper_params):
        super(GRUAdversary, self).__init__()

        self.hyper_params = hyper_params
        self.gru = nn.GRU(
            input_size=hyper_params['item_embed_size'] +
            hyper_params['latent_size'],
            hidden_size=128,
            batch_first=True
        )
        self.linear = nn.Linear(128, 1)

    def forward(self, x, z, padding):
        x = F.gelu(self.gru(torch.cat([x, z], dim=-1))[0])
        x = self.linear(x).squeeze(2)
        return (1.0 - padding.float()) * x
